package com.bjb.vr.manager.filter;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.SerializerProvider;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringEscapeUtils;
import org.apache.commons.lang3.StringUtils;
import org.owasp.validator.html.*;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.io.InputStream;
import java.util.Map;

/**
 * @description: Xss防御拦截器
 * @author: HuJingBo
 * @time: 2022/3/18 13:19
 */
@Slf4j
public class XssRequestWrapper extends HttpServletRequestWrapper {

    private static Policy policy = null;

    static {
        try {
            // 获取策略文件路径，策略文件需要放到项目的classpath下
            InputStream stream = XssRequestWrapper.class.getClassLoader().getResourceAsStream("antisamy-ebay.xml");
            // 指定策略文件
            policy = Policy.getInstance(stream);
        } catch (PolicyException e) {
            e.printStackTrace();
        }
    }

    public XssRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    /**
     * 过滤请求头
     *
     * @param name 参数名
     * @return 参数值
     */
    @Override
    public String getHeader(String name) {
        String header = super.getHeader(name);
        // 如果Header为空，则直接返回，否则进行清洗
        return StringUtils.isBlank(header) ? header : xssClean(header);
    }

    /**
     * 过滤请求参数
     *
     * @param name 参数名
     * @return 参数值
     */
    @Override
    public String getParameter(String name) {
        String parameter = super.getParameter(name);
        // 如果Parameter为空，则直接返回，否则进行清洗
        return StringUtils.isBlank(parameter) ? parameter : xssClean(parameter);
    }

    /**
     * 过滤请求参数（一个参数可以有多个值）
     *
     * @param name 参数名
     * @return 参数值数组
     */
    @Override
    public String[] getParameterValues(String name) {
        String[] parameterValues = super.getParameterValues(name);
        if (parameterValues != null) {
            int length = parameterValues.length;
            String[] newParameterValues = new String[length];
            for (int i = 0; i < length; i++) {
                newParameterValues[i] = xssClean(parameterValues[i]);
            }
            return newParameterValues;
        }
        return super.getParameterValues(name);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> requestMap = super.getParameterMap();
        requestMap.forEach((key, value) -> {
            for (int i = 0; i < value.length; i++) {
                log.info(value[i]);
                value[i] = xssClean(value[i]);
                log.info(value[i]);
            }
        });
        return requestMap;
    }

    /**
     * 使用AntiSamy清洗数据
     *
     * @param value 需要清洗的数据
     * @return 清洗后的数据
     */
    private String xssClean(String value) {
        try {
            AntiSamy antiSamy = new AntiSamy();
            // 使用AntiSamy清洗数据
            final CleanResults cleanResults = antiSamy.scan(value, policy);
            // 获得安全的HTML输出
            value = cleanResults.getCleanHTML();
            // 对转义的HTML特殊字符（<、>、"等）进行反转义，因为AntiSamy调用scan方法时会将特殊字符转义
            value = StringEscapeUtils.unescapeHtml4(value);
            // AntiSamy会把“&nbsp;”转换成乱码，把双引号转换成"&quot;" 先将&nbsp;的乱码替换为空，双引号的乱码替换为双引号
            value = value.replaceAll(antiSamy.scan("&nbsp;", policy).getCleanHTML(), "");
            value = value.replaceAll(antiSamy.scan("\"", policy).getCleanHTML(), "\"");
            return value;
        } catch (ScanException | PolicyException e) {
            log.error("使用AntiSamy清洗数据时发生异常>>> ", e);
        }
        return value;
    }

    /**
     * 通过修改Json序列化的方式来完成Json格式的XSS过滤
     */
    public static class XssStringJsonSerializer extends JsonSerializer<String> {

        @Override
        public Class<String> handledType() {
            return String.class;
        }

        @Override
        public void serialize(String value, JsonGenerator gen, SerializerProvider serializers) throws IOException {
            if (StringUtils.isNotBlank(value) && !"Not Found".equals(value)) {
                try {
                    AntiSamy antiSamy = new AntiSamy();
                    final CleanResults cleanResults = antiSamy.scan(value, policy);
                    gen.writeString(StringEscapeUtils.unescapeHtml4(cleanResults.getCleanHTML()));
                } catch (ScanException | PolicyException e) {
                    e.printStackTrace();
                }
            } else {
                gen.writeString(value);
            }
        }
    }
}